%%--------------------------------------------------------------------
%% Copyright (c) 2025 EMQ Technologies Co., Ltd. All Rights Reserved.
%%--------------------------------------------------------------------
-module(emqx_ds_optimistic_tx).
-moduledoc """
This module implements a per-shard singleton process
facilitating optimistic transactions. This process is dynamically
spawned on the leader node.

Optimistic transaction leader tracks recent writes and verifies
transactions from the clients. Transactions containing reads that
may race with the recently updated topics are rejected. Valid
transactions are "cooked" and added to the buffer.

Buffer is periodically flushed in a single call to the backend.

Leader is also tasked with keeping and incrementing the transaction
serial for the pending transactions. The latest serial is committed
to the storage together with the batches.

Potential split brain situations are handled optimistically: the
backend can reject flush request.
""".

%% API:
-export([
    ls/0,
    config_change/1,
    start_link/3,

    where/2,

    dirty_append/3,
    add_generation/1,

    new_kv_tx_ctx/5,
    commit_kv_tx/4,
    tx_commit_outcome/1,

    get_monotonic_timestamp/0
]).

%% Internal exports
-export([
    init/4,
    code_change/2,
    loop/1
]).

%% System exports
-export([
    system_continue/3,
    system_terminate/4,
    system_code_change/4
]).

-export_type([
    tx/0,
    ctx/0,
    serial/0,
    runtime_config/0,
    batch/0
]).

-include("emqx_ds.hrl").
-include("emqx_ds_builtin_tx.hrl").
-include("emqx_ds_optimistic_tx_internals.hrl").
-include_lib("snabbkaffe/include/snabbkaffe.hrl").

-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-endif.

%%================================================================================
%% Type declarations
%%================================================================================

-type tx() :: #ds_tx{}.

-type serial() :: integer().

-type batch() :: [{emqx_ds:generation(), [_CookedTx]}].

%% Get the last committed transaction ID on the local replica
-callback otx_get_tx_serial(emqx_ds:db(), emqx_ds:shard()) ->
    {ok, serial()} | undefined.

%% Cook a raw transaction into a data structure that will be stored in
%% the buffer until flush.
-callback otx_prepare_tx(
    {emqx_ds:db(), emqx_ds:shard()},
    emqx_ds:generation(),
    _SerialBin :: binary(),
    emqx_ds:tx_ops(),
    _MiscOpts :: map()
) ->
    {ok, _CookedTx} | emqx_ds:error(_).

%% Commit a batch of cooked transactions to the storage.
-callback otx_commit_tx_batch(
    {emqx_ds:db(), emqx_ds:shard()},
    _OldSerial :: serial(),
    _NewSerial :: serial(),
    _NewTimestamp :: emqx_ds:time(),
    batch()
) ->
    ok | emqx_ds:error(_).

-callback otx_add_generation(emqx_ds:db(), emqx_ds:shard(), emqx_ds:time()) ->
    ok | emqx_ds:error(_).

%% Lookup a value identified by topic and timestamp. This callback is
%% executed only on the leader node.
-callback otx_lookup_ttv(
    {emqx_ds:db(), emqx_ds:shard()},
    emqx_ds:generation(),
    emqx_ds:topic(),
    emqx_ds:time()
) ->
    {ok, binary()} | undefined | emqx_ds:error(_).

%% Configuration that can be changed in the runtime.
-type runtime_config() :: #{
    flush_interval := non_neg_integer(),
    idle_flush_interval := non_neg_integer(),
    max_items := pos_integer(),
    conflict_window := pos_integer()
}.

%% Executed by the leader when it updates its runtime config
-callback otx_get_runtime_config(emqx_ds:db()) -> runtime_config().

%% Executed in the leader process when it attempts to step up.
-callback otx_become_leader(emqx_ds:db(), emqx_ds:shard()) ->
    {ok, serial(), emqx_ds:time()} | {error, _}.

%% Executed by the leader
-callback otx_get_latest_generation(emqx_ds:db(), emqx_ds:shard()) -> emqx_ds:generation().

%% Executed by the readers
-callback otx_get_leader(emqx_ds:db(), emqx_ds:shard()) -> pid() | undefined.

%% Executed by the leader prior to flush. `true': ok, leader can write
%% more data, `false': quota exceeded. Input parameter is the
%% `db_group' of the DB.
-callback otx_check_soft_quota(_DBGroup) -> boolean().

-type ctx() :: #kv_tx_ctx{}.
-type process_ref() :: pid() | atom() | {atom(), node()}.
-type leader() :: process_ref() | undefined.

%% Timeouts
%%   Flush pending transactions to the storage:
-define(timeout_flush, timeout_flush).
-define(timeout_quota_check, timeout_quota_check).

-record(gen_data, {
    dirty_w :: emqx_ds_tx_conflict_trie:t(),
    dirty_d :: emqx_ds_tx_conflict_trie:t(),
    buffer = [],
    pending_replies = []
}).

-type gen_data() :: #gen_data{}.

-type generations() :: #{emqx_ds:generation() => gen_data()}.

-record(d, {
    parent :: pid(),
    version :: string(),
    db :: emqx_ds:db(),
    shard :: emqx_ds:shard(),
    cbm :: module(),
    metrics :: emqx_ds_builtin_metrics:shard_metrics_id(),
    %% Timestamp when the FSM changed state from idle to pending, used
    %% to estimate latency induced by buffering:
    entered_pending_at :: integer() | undefined,
    flush_interval :: pos_integer(),
    idle_flush_interval :: pos_integer(),
    rotate_interval :: pos_integer(),
    announce_interval :: pos_integer(),
    max_items :: pos_integer(),
    last_rotate_ts,
    n_items = 0,
    serial :: non_neg_integer() | undefined,
    timestamp :: emqx_ds:time() | undefined,
    committed_serial :: non_neg_integer() | undefined,
    pending_dirty = [] :: [dirty_append()],
    gens = #{} :: generations(),
    %% `true' = writes are allowed, `false' = writes are not allowed
    is_within_quota :: boolean() | undefined,
    %% This timer activates while `is_within_quota' = `false'. This
    %% timer is used to periodacally re-check the quota:
    quota_check_timer :: reference() | undefined
}).
-type d() :: #d{}.

-define(ds_tx_monotonic_ts, ds_tx_monotonic_ts).

%% Entry that is stored in the process dictionary of a process that
%% initiated commit of the transaction:
-define(pending_commit_timer, emqx_ds_optimistic_tx_pending_commit).

-record(cast_dirty_append, {
    reply_to :: pid() | reference() | undefined,
    ref :: reference() | undefined,
    data :: emqx_ds:dirty_append_data(),
    reserved
}).
-type dirty_append() :: #cast_dirty_append{}.

-record(call_add_generation, {}).

%%================================================================================
%% API functions
%%================================================================================

-spec where(emqx_ds:db(), emqx_ds:shard()) -> pid() | undefined.
where(DB, Shard) ->
    gproc:where(?name(DB, Shard)).

-spec start_link(emqx_ds:db(), emqx_ds:shard(), module()) -> {ok, pid()}.
start_link(DB, Shard, CBM) ->
    case where(DB, Shard) of
        undefined ->
            proc_lib:start_link(emqx_ds_optimistic_tx_wrapper, init, [self(), DB, Shard, CBM]);
        Pid when is_pid(Pid) ->
            {error, {already_started, Pid}}
    end.

-spec ls() -> [{emqx_ds:db(), emqx_ds:shard()}].
ls() ->
    MS = {{?name('$1', '$2'), '_', '_'}, [], [{{'$1', '$2'}}]},
    gproc:select({local, names}, [MS]).

-doc """
Notify local workers that DB config has changed.

The backends can call this function from `CBM:update_db_config`.
""".
-spec config_change(emqx_ds:db()) -> ok.
config_change(DB) ->
    MS = {{?name(DB, '_'), '$1', '_'}, [], ['$1']},
    lists:foreach(
        fun(Server) ->
            cast(Server, config_change)
        end,
        gproc:select({local, names}, [MS])
    ).

-spec dirty_append(leader(), emqx_ds:dirty_append_opts(), emqx_ds:dirty_append_data()) ->
    reference() | noreply.
dirty_append(Leader, Opts, Data = [{_, ?ds_tx_ts_monotonic, _} | _]) ->
    Reply = maps:get(reply, Opts, true),
    case Leader of
        P when P =/= undefined ->
            case Reply of
                true ->
                    Alias = Result = monitor(process, Leader, [{alias, reply_demonitor}]);
                false ->
                    Alias = undefined,
                    Result = noreply
            end,
            cast(Leader, #cast_dirty_append{
                reply_to = Alias,
                ref = Alias,
                data = Data
            }),
            Result;
        undefined when Reply ->
            Ref = make_ref(),
            reply_error(self(), Ref, dirty, recoverable, leader_down),
            Ref;
        undefined ->
            noreply
    end.

-spec add_generation(leader()) -> ok | emqx_ds:error(_).
add_generation(Leader) ->
    case Leader of
        P when P =/= undefined ->
            call(Leader, #call_add_generation{}, infinity);
        undefined ->
            ?err_rec(leader_down)
    end.

-spec new_kv_tx_ctx(
    module(),
    emqx_ds:db(),
    emqx_ds:shard(),
    emqx_ds:generation() | latest,
    emqx_ds:transaction_opts()
) -> {ok, ctx()} | emqx_ds:error(_).
new_kv_tx_ctx(CBM, DB, Shard, Generation0, Opts) ->
    case Generation0 of
        latest ->
            Generation = CBM:otx_get_latest_generation(DB, Shard),
            IsLatest = true;
        Generation when is_integer(Generation) ->
            IsLatest = false
    end,
    maybe
        Leader = CBM:otx_get_leader(DB, Shard),
        true ?= is_pid(Leader),
        {ok, Serial} ?= CBM:otx_get_tx_serial(DB, Shard),
        Ctx = #kv_tx_ctx{
            shard = Shard,
            leader = Leader,
            serial = Serial,
            generation = Generation,
            latest_generation = IsLatest,
            opts = Opts
        },
        {ok, Ctx}
    else
        _ ->
            ?err_rec(leader_down)
    end.

-spec commit_kv_tx(leader(), emqx_ds:db(), ctx(), emqx_ds:tx_ops()) -> reference().
commit_kv_tx(Leader, DB, Ctx = #kv_tx_ctx{opts = #{timeout := Timeout}}, Ops) ->
    ?tp(emqx_ds_optimistic_tx_commit_begin, #{db => DB, ctx => Ctx, ops => Ops}),
    case Leader of
        P when P =/= undefined ->
            Alias = monitor(process, Leader, [{alias, reply_demonitor}]),
            TRef = emqx_ds_lib:send_after(Timeout, self(), tx_timeout_msg(Alias)),
            put({?pending_commit_timer, Alias}, TRef),
            cast(
                Leader,
                #ds_tx{ctx = Ctx, ops = Ops, from = Alias, ref = Alias}
            ),
            Alias;
        undefined ->
            Alias = make_ref(),
            self() ! ?ds_tx_commit_error(Alias, undefined, recoverable, leader_unavailable),
            Alias
    end.

-spec tx_commit_outcome(term()) -> {ok, emqx_ds:tx_serial()} | emqx_ds:error(_).
tx_commit_outcome(Reply) ->
    case Reply of
        ?ds_tx_commit_ok(Ref, _Meta, Serial) ->
            client_post_commit_cleanup(Ref),
            {ok, Serial};
        ?ds_tx_commit_error(Ref, _Meta, Class, Info) ->
            client_post_commit_cleanup(Ref),
            {error, Class, Info};
        {'DOWN', Ref, Type, Object, Info} ->
            client_post_commit_cleanup(Ref),
            ?err_rec({Type, Object, Info})
    end.

-doc """
Get shard's monotonic timestamp.

This function must be called only in the optimistic_tx process, inside
`with_time_and_generation` wrapper.
""".
-spec get_monotonic_timestamp() -> emqx_ds:time().
get_monotonic_timestamp() ->
    TS = max(get(?ds_tx_monotonic_ts) + 1, erlang:system_time(microsecond)),
    put(?ds_tx_monotonic_ts, TS),
    TS.

%%================================================================================
%% Internal exports
%%================================================================================

-spec init(pid(), emqx_ds:db(), emqx_ds:shard(), module()) -> no_return().
init(Parent, DB, Shard, CBM) ->
    erlang:process_flag(trap_exit, true),
    logger:update_process_metadata(#{db => DB, shard => Shard}),
    ?tp(info, ds_otx_start, #{db => DB, shard => Shard}),
    #{
        flush_interval := FI,
        idle_flush_interval := IFI,
        conflict_window := CW,
        max_items := MaxItems
    } = CBM:otx_get_runtime_config(DB),
    ok = emqx_ds_builtin_metrics:init_for_buffer(DB, Shard),
    D = #d{
        parent = Parent,
        version = version(),
        db = DB,
        shard = Shard,
        cbm = CBM,
        metrics = emqx_ds_builtin_metrics:metric_id([{db, DB}, {shard, Shard}]),
        rotate_interval = CW,
        last_rotate_ts = erlang:monotonic_time(millisecond),
        flush_interval = FI,
        idle_flush_interval = IFI,
        announce_interval = 5_000,
        max_items = MaxItems
    },
    proc_lib:init_ack(Parent, {ok, self()}),
    %% Async init:
    maybe
        {ok, Serial, Timestamp} ?= CBM:otx_become_leader(DB, Shard),
        %% Issue a dummy transaction to trigger metadata update:
        ok ?= CBM:otx_commit_tx_batch({DB, Shard}, Serial, Serial, Timestamp, []),
        ?tp(info, ds_otx_up, #{serial => Serial, db => DB, shard => Shard, ts => Timestamp}),
        loop(
            check_soft_quota(D#d{
                serial = Serial,
                timestamp = Timestamp,
                committed_serial = Serial
            })
        )
    else
        Err ->
            terminate({init_failed, Err}, D)
    end.

-spec loop(d()) -> no_return().
loop(D0) ->
    #d{
        db = DB,
        shard = Shard,
        idle_flush_interval = IdleFlushInterval,
        n_items = NItems
    } = D0,
    Sleep =
        case NItems of
            0 -> infinity;
            _ -> IdleFlushInterval
        end,
    receive
        {'$gen_cast', Tx = #ds_tx{}} ->
            case handle_tx(check_soft_quota(D0), Tx) of
                {ok, _PresumedCommitSerial, D} ->
                    loop(finalize_add_pending(D));
                aborted ->
                    loop(D0)
            end;
        {'$gen_cast', #cast_dirty_append{} = Tx} ->
            D = handle_dirty_append(Tx, check_soft_quota(D0)),
            loop(finalize_add_pending(D));
        {'$gen_call', From, #call_add_generation{}} ->
            {Reply, D} = handle_add_generation(D0),
            gen:reply(From, Reply),
            loop(D);
        {'$gen_cast', config_change} ->
            loop(do_update_config(D0));
        ?timeout_flush ->
            loop(flush(D0));
        ?timeout_quota_check ->
            D = D0#d{quota_check_timer = undefined},
            loop(check_soft_quota(D));
        {'EXIT', _From, Reason} ->
            case Reason of
                normal -> loop(D0);
                _ -> terminate(Reason, D0)
            end;
        {system, From, SysReq} ->
            Debug = [],
            sys:handle_system_msg(SysReq, From, D0#d.parent, ?MODULE, Debug, D0);
        {'$gen_call', From, Other} ->
            ?tp(warning, emqx_ds_optimistic_tx_unknown_call, #{
                call => Other,
                db => DB,
                shard => Shard,
                from => From
            }),
            gen:reply(From, ?err_unrec({unknown_call, Other})),
            loop(D0);
        Other ->
            ?tp(warning, emqx_ds_optimistic_tx_unknown_event, #{
                event => Other,
                db => DB,
                shard => Shard
            }),
            loop(D0)
    after Sleep ->
        loop(flush(D0))
    end.

code_change(D, _Extra) ->
    D#d{version = version()}.

%%================================================================================
%% System functions
%%================================================================================

system_continue(Parent, _Debug, D0) ->
    ?MODULE:loop(D0#d{parent = Parent}).

-spec system_terminate(_, _, _, _) -> no_return().
system_terminate(Reason, _Parent, _Debug, D) ->
    terminate(Reason, D).

system_code_change(D, _Module, _OldVsn, Extra) ->
    {ok, ?MODULE:code_change(D, Extra)}.

%%================================================================================
%% Internal functions
%%================================================================================

-spec terminate(term(), term()) -> no_return().
terminate(Reason, _Data) ->
    Level =
        case Reason =:= normal orelse Reason =:= shutdown of
            true ->
                info;
            false ->
                %% Sleep some to prevent a hot restart loop
                timer:sleep(500),
                error
        end,
    ?tp(Level, ds_otx_terminate, #{reason => Reason}),
    exit(Reason).

-spec flush(d()) -> d().
flush(D0) ->
    %% Execute flush. After flushing the buffer it's safe to rotate
    %% the conflict tree.
    maybe_rotate(do_update_config(do_flush(D0))).

-spec handle_dirty_append(dirty_append(), d()) -> d().
handle_dirty_append(
    #cast_dirty_append{reply_to = From, ref = Ref, reserved = Reserved},
    D = #d{is_within_quota = false}
) ->
    reply_error(From, Ref, Reserved, unrecoverable, storage_quota_exceeded),
    D;
handle_dirty_append(
    DirtyAppend, D = #d{pending_dirty = Buff, n_items = NItems}
) ->
    %% Note: all dirty appends are cooked in a single `otx_prepare_tx` call during the flush:
    D#d{
        pending_dirty = [DirtyAppend | Buff],
        n_items = NItems + 1
    }.

-spec handle_add_generation(d()) -> {Reply, d()} when
    Reply :: ok | emqx_ds:error(_).
handle_add_generation(D = #d{db = DB, shard = Shard, cbm = CBM, timestamp = TS0}) ->
    put(?ds_tx_monotonic_ts, TS0),
    TS = get_monotonic_timestamp(),
    _ = erase(?ds_tx_monotonic_ts),
    Reply = CBM:otx_add_generation(DB, Shard, TS),
    %% TODO: commit an empty batch to persist down timestamp:
    {Reply, D#d{timestamp = TS}}.

-spec send_dirty_append_replies(ok | false | emqx_ds:error(_), serial(), [dirty_append()]) -> ok.
send_dirty_append_replies(Result, Serial, PendingReplies) ->
    case Result of
        ok ->
            lists:foreach(
                fun(#cast_dirty_append{reply_to = From, ref = Ref, reserved = Reserved}) ->
                    reply_success(From, Ref, Reserved, Serial)
                end,
                PendingReplies
            );
        Other ->
            case Other of
                false ->
                    %% Failed to cook:
                    ErrClass = recoverable,
                    Error = aborted;
                {error, ErrClass, Error} ->
                    ok
            end,
            lists:foreach(
                fun(#cast_dirty_append{reply_to = From, ref = Ref, reserved = Reserved}) ->
                    reply_error(From, Ref, Reserved, ErrClass, Error)
                end,
                PendingReplies
            )
    end.

-spec cook_dirty_appends(d()) -> {Success, serial() | undefined, [dirty_append()], d()} when
    Success :: boolean().
cook_dirty_appends(D = #d{pending_dirty = []}) ->
    {true, undefined, [], D};
cook_dirty_appends(D0 = #d{db = DB, shard = Shard, cbm = CBM, pending_dirty = Pending}) ->
    Gen = CBM:otx_get_latest_generation(DB, Shard),
    Result = with_time_and_generation(
        Gen,
        D0,
        fun(GenData, PresumedCommitSerial) ->
            do_cook_dirty_appends(DB, Shard, CBM, Gen, GenData, PresumedCommitSerial, Pending)
        end
    ),
    %% Optimization: here we return the original `Pending' list to be
    %% passed to `send_dirty_append_replies' function. This should
    %% prevent extra list allocations when sending replies.
    case Result of
        {ok, PresumedCommitSerial, D} ->
            {true, PresumedCommitSerial, Pending, D#d{pending_dirty = []}};
        aborted ->
            {false, undefined, Pending, D0#d{pending_dirty = []}}
    end.

do_cook_dirty_appends(DB, Shard, CBM, Gen, GenData, PresumedCommitSerial, Pending) ->
    maybe
        #gen_data{
            buffer = Buff,
            pending_replies = _Pending0
        } = GenData,
        Writes = lists:foldl(
            fun(#cast_dirty_append{data = Data}, Acc) ->
                Data ++ Acc
            end,
            [],
            Pending
        ),
        {ok, CookedDirty} ?=
            prepare_tx(
                CBM,
                {DB, Shard},
                Gen,
                serial_bin(PresumedCommitSerial),
                #{?ds_tx_write => Writes},
                #{}
            ),
        {ok, GenData#gen_data{
            buffer = [CookedDirty | Buff]
        }}
    else
        _ -> aborted
    end.

-doc """
This function must be called after adding a new item pending for commit.
It (re)schedules flushing of the data.
""".
-spec finalize_add_pending(d()) -> d().
finalize_add_pending(
    D0 = #d{
        entered_pending_at = EP, flush_interval = FlushInterval, n_items = N, max_items = MaxItems
    }
) ->
    D =
        case EP of
            _ when is_integer(EP) ->
                D0;
            undefined ->
                %% Schedule unconditional flush after FlushInterval:
                erlang:send_after(FlushInterval, self(), ?timeout_flush),
                D0#d{entered_pending_at = erlang:monotonic_time(millisecond)}
        end,
    case N >= MaxItems andalso MaxItems > 0 of
        true ->
            flush(D);
        false ->
            D
    end.

-doc """
A wrapper that sets up the environment for calling CBM:tx_prepare_tx.

If successful, it updates generation data, current monotonic time and
commit serial in the global state.
""".
-spec with_time_and_generation(emqx_ds:generation(), d(), fun(
    (gen_data(), serial()) -> {ok, gen_data()} | aborted
)) ->
    {ok, serial(), d()} | aborted.
with_time_and_generation(
    Gen,
    D = #d{
        n_items = NItems,
        serial = Serial,
        timestamp = Timestamp,
        gens = Gens
    },
    Fun
) ->
    case Gens of
        #{Gen := GenData0} ->
            ok;
        #{} ->
            GenData0 = #gen_data{
                dirty_w = emqx_ds_tx_conflict_trie:new(Serial, infinity),
                dirty_d = emqx_ds_tx_conflict_trie:new(Serial, infinity)
            }
    end,
    PresumedCommitSerial = Serial + 1,
    put(?ds_tx_monotonic_ts, Timestamp),
    case Fun(GenData0, PresumedCommitSerial) of
        {ok, GenData} ->
            {ok, PresumedCommitSerial, D#d{
                serial = PresumedCommitSerial,
                gens = Gens#{Gen => GenData},
                timestamp = erase(?ds_tx_monotonic_ts),
                n_items = NItems + 1
            }};
        aborted ->
            _ = erase(?ds_tx_monotonic_ts),
            aborted
    end.

-spec handle_tx(d(), tx()) -> {ok, serial(), d()} | aborted.
handle_tx(
    D = #d{
        db = DB,
        shard = Shard,
        cbm = CBM,
        serial = Serial,
        committed_serial = SafeToReadSerial,
        is_within_quota = WithinQuota
    },
    Tx
) ->
    #ds_tx{
        ref = TxRef,
        ctx = #kv_tx_ctx{generation = Gen, latest_generation = IsLatest}
    } = Tx,
    with_time_and_generation(
        Gen,
        D,
        fun(GenData, PresumedCommitSerial) ->
            DBShard = {DB, Shard},
            ?tp(
                emqx_ds_optimistic_tx_commit_received,
                #{
                    shard => DBShard,
                    serial => Serial,
                    committed_serial => SafeToReadSerial,
                    ref => TxRef,
                    tx => Tx,
                    generation => {IsLatest, Gen}
                }
            ),
            try_schedule_transaction(
                DBShard, Gen, CBM, SafeToReadSerial, PresumedCommitSerial, WithinQuota, Tx, GenData
            )
        end
    ).

-doc """
This function verifies that a transaction doesn't conflict with data
that was either committed or is scheduled for commit, and that
transaction preconditions are satisfied.

If all checks pass then the transaction is cooked and scheduled for
commit, and conflict tracking is updated.
""".
-spec try_schedule_transaction(
    {emqx_ds:db(), emqx_ds:shard()},
    emqx_ds:generation(),
    module(),
    serial(),
    serial(),
    boolean(),
    tx(),
    gen_data()
) ->
    {ok, gen_data()} | aborted.
try_schedule_transaction(
    DBShard,
    Gen,
    CBM,
    SafeToReadSerial,
    PresumedCommitSerial,
    WithinQuota,
    Tx,
    GS = #gen_data{dirty_w = DirtyW0, dirty_d = DirtyD0, buffer = Buff, pending_replies = Pending}
) ->
    #ds_tx{
        ctx = Ctx,
        ops = Ops0,
        from = From,
        ref = Ref,
        meta = Meta
    } = Tx,
    Ops = preprocess_ops(PresumedCommitSerial, Ops0),
    #kv_tx_ctx{serial = TxStartSerial, leader = TxLeader} = Ctx,
    maybe
        ok ?= check_tx_leader(TxLeader),
        ok ?= check_latest_generation(CBM, DBShard, Ctx),
        ok ?= check_quota(WithinQuota, Ops),
        ok ?= check_conflicts(DirtyW0, DirtyD0, TxStartSerial, SafeToReadSerial, Ops),
        ok ?= verify_preconditions(DBShard, CBM, Gen, Ops),
        {ok, CookedTx} ?= prepare_tx(CBM, DBShard, Gen, serial_bin(PresumedCommitSerial), Ops, #{}),
        ?tp(emqx_ds_optimistic_tx_commit_pending, #{ref => Ref}),
        {ok, GS#gen_data{
            dirty_w = update_dirty_w(PresumedCommitSerial, Ops, DirtyW0),
            dirty_d = update_dirty_d(PresumedCommitSerial, Ops, DirtyD0),
            buffer = [CookedTx | Buff],
            pending_replies = [{From, Ref, Meta, PresumedCommitSerial} | Pending]
        }}
    else
        {error, Class, Error} ->
            ?tp(
                debug,
                emqx_ds_optimistic_tx_commit_abort,
                #{
                    ref => Ref, ec => Class, reason => Error, stage => prepare
                }
            ),
            reply_error(From, Ref, Meta, Class, Error),
            aborted
    end.

-doc """
Check that the transaction was created while this process was the leader.
""".
-spec check_tx_leader(pid()) -> ok | emqx_ds:error(_).
check_tx_leader(TxLeader) ->
    case self() of
        TxLeader ->
            ok;
        _ ->
            ?err_rec({leader_conflict, TxLeader, self()})
    end.

-spec check_latest_generation(module(), {emqx_ds:db(), emqx_ds:shard()}, ctx()) ->
    ok | emqx_ds:error(_).
check_latest_generation(_, _, #kv_tx_ctx{latest_generation = false}) ->
    %% This transaction was created for a constant generation that
    %% doesn't have to be the latest. This always succeeds.
    ok;
check_latest_generation(CBM, {DB, Shard}, #kv_tx_ctx{latest_generation = true, generation = Gen}) ->
    case CBM:otx_get_latest_generation(DB, Shard) of
        Gen ->
            ok;
        Latest ->
            ?err_rec({not_the_latest_generation, Gen, Latest})
    end.

check_quota(true, _) ->
    ok;
check_quota(false, Ops) ->
    case maps:get(?ds_tx_write, Ops, []) of
        [] ->
            ok;
        _ ->
            ?err_unrec(storage_quota_exceeded)
    end.

-spec update_dirty_d(serial(), emqx_ds:tx_ops(), emqx_ds_tx_conflict_trie:t()) ->
    emqx_ds_tx_conflict_trie:t().
update_dirty_d(Serial, Ops, Dirty) ->
    %% Mark all deleted topics as dirty:
    lists:foldl(
        fun({TF, FromTime, ToTime}, Acc) ->
            emqx_ds_tx_conflict_trie:push_topic(
                emqx_ds_tx_conflict_trie:topic_filter_to_conflict_domain(TF),
                FromTime,
                ToTime,
                Serial,
                Acc
            )
        end,
        Dirty,
        maps:get(?ds_tx_delete_topic, Ops, [])
    ).

-spec update_dirty_w(serial(), emqx_ds:tx_ops(), emqx_ds_tx_conflict_trie:t()) ->
    emqx_ds_tx_conflict_trie:t().
update_dirty_w(Serial, Ops, Dirty) ->
    %% Mark written topics as dirty:
    lists:foldl(
        fun({TF, TS, _Val}, Acc) ->
            emqx_ds_tx_conflict_trie:push_topic(TF, TS, TS, Serial, Acc)
        end,
        Dirty,
        maps:get(?ds_tx_write, Ops, [])
    ).

-spec do_flush(d()) -> d().
do_flush(D = #d{n_items = 0}) ->
    D;
do_flush(D0 = #d{n_items = NItems}) ->
    %% Note: we take n_items before cooking_dirty_append to avoid
    %% off-by-one error when it adds a batch.
    {CookDirtySuccess, PresumedDirtyCommitSerial, OldDirtyAppends, D} = cook_dirty_appends(D0),
    #d{
        db = DB,
        shard = Shard,
        cbm = CBM,
        metrics = Metrics,
        entered_pending_at = EnteredPendingAt,
        committed_serial = SerCtl,
        serial = Serial,
        gens = Gens0,
        timestamp = Timestamp
    } = D,
    DBShard = {DB, Shard},
    Batch = make_batch(Gens0),
    T0 = erlang:monotonic_time(microsecond),
    Result = CBM:otx_commit_tx_batch(DBShard, SerCtl, Serial, Timestamp, Batch),
    T1 = erlang:monotonic_time(microsecond),
    emqx_ds_builtin_metrics:observe_buffer_flush_time(Metrics, T1 - T0),
    Gens = clean_buffers_and_reply(Result, Gens0),
    send_dirty_append_replies(
        CookDirtySuccess andalso Result, PresumedDirtyCommitSerial, OldDirtyAppends
    ),
    case Result of
        ok ->
            Latency = erlang:monotonic_time(millisecond) - EnteredPendingAt,
            emqx_ds_builtin_metrics:inc_buffer_batches(Metrics),
            emqx_ds_builtin_metrics:inc_buffer_messages(Metrics, NItems),
            emqx_ds_builtin_metrics:observe_buffer_latency(Metrics, Latency),
            D#d{
                entered_pending_at = undefined,
                committed_serial = Serial,
                gens = Gens,
                n_items = 0
            };
        _ ->
            emqx_ds_builtin_metrics:inc_buffer_batches_failed(Metrics),
            exit({flush_failed, #{db => DB, shard => Shard, result => Result}})
    end.

-spec do_update_config(d()) -> d().
do_update_config(D = #d{db = DB, cbm = CBM}) ->
    #{
        flush_interval := FI,
        idle_flush_interval := IFI,
        max_items := MaxItems,
        conflict_window := CW
    } = CBM:otx_get_runtime_config(DB),
    D#d{
        flush_interval = FI,
        idle_flush_interval = IFI,
        max_items = MaxItems,
        rotate_interval = CW
    }.

-spec make_batch(generations()) -> batch().
make_batch(Generations) ->
    maps:fold(
        fun(Generation, #gen_data{buffer = Buf}, Acc) ->
            case Buf of
                [] ->
                    %% Avoid touching generations that don't have
                    %% data. They could be deleted. TODO: drop deleted
                    %% generations from the state record.
                    Acc;
                _ ->
                    [{Generation, Buf} | Acc]
            end
        end,
        [],
        Generations
    ).

-spec clean_buffers_and_reply(ok | emqx_ds:error(_), generations()) -> generations().
clean_buffers_and_reply(Reply, Generations) ->
    maps:map(
        fun(_, GenData = #gen_data{pending_replies = Waiting}) ->
            _ =
                case Reply of
                    ok ->
                        [
                            reply_success(From, Ref, Meta, CommitSerial)
                         || {From, Ref, Meta, CommitSerial} <- Waiting
                        ];
                    {error, Class, Err} ->
                        [
                            reply_error(From, Ref, Meta, Class, Err)
                         || {From, Ref, Meta, _CommitSerial} <- Waiting
                        ]
                end,
            GenData#gen_data{
                buffer = [],
                pending_replies = []
            }
        end,
        Generations
    ).

-spec serial_bin(serial()) -> binary().
serial_bin(A) ->
    <<A:128>>.

-spec maybe_rotate(d()) -> d().
maybe_rotate(D = #d{rotate_interval = RI, last_rotate_ts = LastRotTS}) ->
    case erlang:monotonic_time(millisecond) - LastRotTS > RI of
        true ->
            rotate(D);
        false ->
            D
    end.

-spec rotate(d()) -> d().
rotate(D = #d{gens = Gens0}) ->
    Gens = maps:map(
        fun(_, GS = #gen_data{dirty_w = Dirty}) ->
            GS#gen_data{
                dirty_w = emqx_ds_tx_conflict_trie:rotate(Dirty)
            }
        end,
        Gens0
    ),
    D#d{
        gens = Gens,
        last_rotate_ts = erlang:monotonic_time(millisecond)
    }.

-spec preprocess_ops(serial(), emqx_ds:tx_ops()) -> emqx_ds:tx_ops().
preprocess_ops(_PresumedCommitSerial, Tx) ->
    MonotonicTS = get_monotonic_timestamp(),
    maps:map(
        fun
            (?ds_tx_delete_topic, Ops) ->
                preprocess_deletes(MonotonicTS, Ops);
            (_Key, Ops) ->
                Ops
        end,
        Tx
    ).

-spec preprocess_deletes(emqx_ds:time(), [emqx_ds:topic_range()]) ->
    [{emqx_ds:topic(), emqx_ds:time(), emqx_ds:time() | infinity}].
preprocess_deletes(MonotonicTS, Ops) ->
    lists:map(
        fun
            ({TopicFilter, From, ?ds_tx_ts_monotonic}) ->
                {TopicFilter, From, MonotonicTS};
            (A) ->
                A
        end,
        Ops
    ).

-spec check_conflicts(
    emqx_ds_tx_conflict_trie:t(), emqx_ds_tx_conflict_trie:t(), serial(), serial(), emqx_ds:tx_ops()
) ->
    ok | emqx_ds:error(_).
check_conflicts(DirtyW, DirtyD, TxStartSerial, SafeToReadSerial, Ops) ->
    maybe
        %% Check reads against writes and deletions:
        Reads = maps:get(?ds_tx_read, Ops, []),
        ok ?= do_check_conflicts(DirtyW, TxStartSerial, Reads),
        ok ?= do_check_conflicts(DirtyD, TxStartSerial, Reads),
        %% Deletion of topics involves scanning of the storage. We
        %% can't do it when there is potentially buffered data:
        ok ?= do_check_conflicts(DirtyW, SafeToReadSerial, maps:get(?ds_tx_delete_topic, Ops, [])),
        %% Verifying precondition requires reading from the storage
        %% too. Make sure we don't ignore the cache:
        ExpectedTopics = [{Topic, TS, TS} || {Topic, TS, _} <- maps:get(?ds_tx_expected, Ops, [])],
        ok ?= do_check_conflicts(DirtyW, SafeToReadSerial, ExpectedTopics),
        ok ?= do_check_conflicts(DirtyD, SafeToReadSerial, ExpectedTopics),
        UnexpectedTopics = [{Topic, TS, TS} || {Topic, TS} <- maps:get(?ds_tx_unexpected, Ops, [])],
        ok ?= do_check_conflicts(DirtyW, SafeToReadSerial, UnexpectedTopics)
    end.

-spec do_check_conflicts(emqx_ds_tx_conflict_trie:t(), serial(), [
    {emqx_ds:topic(), emqx_ds:time(), emqx_ds:time() | infinity}
]) -> ok | ?err_rec(_).
do_check_conflicts(Dirty, Serial, Topics) ->
    Errors = lists:filtermap(
        fun({ReadTF, FromTime, ToTime}) ->
            emqx_ds_tx_conflict_trie:is_dirty_topic(
                emqx_ds_tx_conflict_trie:topic_filter_to_conflict_domain(ReadTF),
                FromTime,
                ToTime,
                Serial,
                Dirty
            ) andalso
                {true, {ReadTF, {FromTime, ToTime}, Serial}}
        end,
        Topics
    ),
    case Errors of
        [] ->
            ok;
        _ ->
            ?err_rec({read_conflict, Errors})
    end.

-spec verify_preconditions(
    {emqx_ds:db(), emqx_ds:shard()}, module(), emqx_ds:generation(), emqx_ds:tx_ops()
) ->
    ok | ?err_unrec(_).
verify_preconditions(DBShard, CBM, GenId, Ops) ->
    %% Verify expected values:
    Unrecoverable0 = lists:foldl(
        fun({Topic, Time, ExpectedValue}, Acc) ->
            case CBM:otx_lookup_ttv(DBShard, GenId, Topic, Time) of
                {ok, Value} when
                    ExpectedValue =:= '_';
                    ExpectedValue =:= Value
                ->
                    Acc;
                {ok, Value} ->
                    [#{topic => Topic, expected => ExpectedValue, ts => Time, got => Value} | Acc];
                undefined ->
                    [
                        #{topic => Topic, expected => ExpectedValue, ts => Time, got => undefined}
                        | Acc
                    ]
            end
        end,
        [],
        maps:get(?ds_tx_expected, Ops, [])
    ),
    %% Verify unexpected values:
    Unrecoverable = lists:foldl(
        fun({Topic, Time}, Acc) ->
            case CBM:otx_lookup_ttv(DBShard, GenId, Topic, Time) of
                undefined ->
                    Acc;
                {ok, Value} ->
                    [#{topic => Topic, ts => Time, unexpected => Value} | Acc]
            end
        end,
        Unrecoverable0,
        maps:get(?ds_tx_unexpected, Ops, [])
    ),
    case Unrecoverable of
        [] ->
            ok;
        _ ->
            ?err_unrec({precondition_failed, Unrecoverable})
    end.

tx_timeout_msg(Ref) ->
    ?ds_tx_commit_error(Ref, undefined, unrecoverable, commit_timeout).

reply_success(From, Ref, Meta, Serial) ->
    ?tp(
        emqx_ds_optimistic_tx_commit_success,
        #{client => From, ref => Ref, meta => Meta, serial => Serial}
    ),
    case From of
        _ when is_pid(From) orelse is_reference(From) ->
            From ! ?ds_tx_commit_ok(Ref, Meta, serial_bin(Serial)),
            ok;
        undefined ->
            ok
    end.

reply_error(From, Ref, Meta, Class, Error) ->
    ?tp(
        debug,
        emqx_ds_optimistic_tx_commit_abort,
        #{client => From, ref => Ref, meta => Meta, ec => Class, reason => Error, stage => final}
    ),
    case From of
        _ when is_pid(From) orelse is_reference(From) ->
            From ! ?ds_tx_commit_error(Ref, Meta, Class, Error),
            ok;
        undefined ->
            ok
    end.

-spec client_post_commit_cleanup(reference()) -> ok.
client_post_commit_cleanup(Alias) ->
    %% This cleanup procedure leaves a slight chance that timeout
    %% message will still arrive. However, trying to flush it is a bad
    %% idea! Erlang cannot use ref trick here to optimize selective
    %% receive, so flush will lead to a full mailbox scan.
    demonitor(Alias),
    case erase({?pending_commit_timer, Alias}) of
        TRef when is_reference(TRef) ->
            _ = erlang:cancel_timer(TRef),
            ok;
        _ ->
            ok
    end.

prepare_tx(CBM, DBShard, Gen, SerialBin, Ops, CookOpts) ->
    try
        CBM:otx_prepare_tx(DBShard, Gen, SerialBin, Ops, CookOpts)
    catch
        EC:Err:Stack ->
            {error, recoverable, {EC, Err, Stack}}
    end.

cast(Server, Message) ->
    try erlang:send(Server, {'$gen_cast', Message}) of
        _ -> ok
    catch
        _:_ -> ok
    end.

call(Server, Request, Timeout) ->
    try gen:call(Server, '$gen_call', Request, Timeout) of
        {ok, Reply} ->
            Reply
    catch
        Class:Err:Stack ->
            erlang:raise(
                Class,
                {Err, {?MODULE, call, [Server, Request, Timeout]}},
                Stack
            )
    end.

version() ->
    {ok, Version} = application:get_key(emqx_durable_storage, vsn),
    Version.

check_soft_quota(
    D = #d{cbm = CBM, db = DB, is_within_quota = OldQuotaState, quota_check_timer = undefined}
) ->
    maybe
        #{runtime := RT} ?= emqx_dsch:get_db_runtime(DB),
        #{db_group := GroupName} ?= RT,
        {ok, Group} ?= emqx_ds:lookup_db_group(GroupName),
        NewQuotaState = CBM:otx_check_soft_quota(Group),
        %% Handle alarms and events:
        case NewQuotaState of
            OldQuotaState ->
                ok;
            false ->
                ?tp(emqx_ds_storage_quota, #{db => DB, exceeded => true}),
                emqx_alarm:safe_activate(
                    quota_alarm(DB),
                    #{},
                    <<"Storage quota has been exceeded for database ", (atom_to_binary(DB))/binary>>
                );
            true ->
                ?tp(emqx_ds_storage_quota, #{db => DB, exceeded => false}),
                emqx_alarm:safe_deactivate(quota_alarm(DB))
        end,
        D#d{
            is_within_quota = NewQuotaState,
            quota_check_timer = maybe_start_quota_check_timer(NewQuotaState)
        }
    else
        Other ->
            ?tp(warning, "invalid_soft_quota_result", #{cbm => CBM, db => DB, result => Other}),
            D
    end;
check_soft_quota(D) ->
    D.

quota_alarm(DB) ->
    <<"db_storage_quota_exceeded:", (atom_to_binary(DB))/binary>>.

maybe_start_quota_check_timer(true) ->
    undefined;
maybe_start_quota_check_timer(false) ->
    Interval = 1000,
    erlang:send_after(Interval, self(), ?timeout_quota_check).

%%================================================================================
%% Tests
%%================================================================================
